#!/usr/bin/env python3
# coding:utf-8

import time
import math
from dataclasses import dataclass, field
from typing import List
import csv

import rclpy
from rclpy.node import Node
from std_msgs.msg import String, Float64, Float64MultiArray, Int64, UInt64MultiArray


@dataclass
class Sub_Topic_Hz:
    name: str
    value_name: str
    fixed_hz: int  # 固有频率
    fixed_hz_diff_time: float  # 固定帧时间差
    msg_type: any = None

    # 记录时间, 数据
    hz: float = 0.0
    min_time: float = 1000.0
    max_time: float = 0.0
    new_start_time: float = 0.0  # 上一秒开始统计时间
    last_time: float = 0.0  # 上一条topic话题时间
    times: List[float] = field(default_factory=list)
    # 统计时间
    accum_dev_time: float = 0.0  # 积累偏差
    std_dev_time: float = 0.0  # 标准偏差


Config = {
    "record_path": "/tmp",
    "csv_record": True,
    "csv_filed": [
        "time",
        "topic",
        "rate",
        "min",
        "max",
        "accum_dev",
        "std_dev",
    ],
}

sub_topic_list = [
    Sub_Topic_Hz("/hz_string", "hz_string", 100, 1.0 / 100, msg_type=String),
    Sub_Topic_Hz("/hz_int", "hz_int", 100, 1.0 / 100, msg_type=Int64),
    Sub_Topic_Hz("/hz_int_array", "hz_int_array", 50, 1.0 / 50, msg_type=UInt64MultiArray),
    Sub_Topic_Hz("/hz_float", "hz_float", 100, 1.0 / 100, msg_type=Float64),
    Sub_Topic_Hz("/hz_float_array", "hz_float_array", 50, 1.0 / 50, msg_type=Float64MultiArray),
]


# import debugpy
# debugpy.listen(6688)
# debugpy.wait_for_client()


class topicSubscribeHz(Node):
    def __init__(self, sub_topic_list, config):
        super().__init__("topic_subscribe_count")
        self.config = config
        self.map_callback = {}
        self.sub_list = []
        for topic in sub_topic_list:
            callback_name = f"callback_{topic.value_name}"
            callback = getattr(self, callback_name)
            sub = self.create_subscription(topic.msg_type, topic.name, callback, 1)
            # 动态添加类成员变量
            setattr(self, topic.value_name, topic)
            self.sub_list.append(sub)
            self.map_callback[callback_name] = callback
        self.init_csv_record()

    def __exit__(self):
        if self.config["csv_record"] == True:
            self.csv_fp.close()

    def init_csv_record(self):
        if self.config["csv_record"] == False:
            return
        time_stamp = time.strftime("%Y%m%d-%H%M%S", time.localtime())
        file = f"{self.config['record_path']}/hz_{time_stamp}.csv"
        fieldnames = self.config["csv_filed"]
        self.csv_fp = open(file, "w")
        self.csv_writer = csv.DictWriter(self.csv_fp, fieldnames=fieldnames)
        self.csv_writer.writeheader()

    def csv_record(self, hz):
        if self.config["csv_record"] == False:
            return
        self.csv_writer.writerow(
            {
                "time": time.strftime("%Y%m%d_%H%M%S", time.localtime()),
                "topic": hz.name,
                "rate": hz.hz,
                "min": hz.min_time,
                "max": hz.max_time,
                "std_dev": hz.std_dev_time,
                "accum_dev": hz.accum_dev_time,
            }
        )

    def cal_hz(self, hz):
        cur_time = time.time()
        diff_time = cur_time - hz.new_start_time
        if hz.new_start_time == 0.0:
            hz.new_start_time = cur_time
        elif diff_time >= 1.0:
            if hz.hz != 0.0:
                hz.accum_dev_time = round(
                    (hz.accum_dev_time + hz.fixed_hz * hz.fixed_hz_diff_time), 3
                )
                hz.min_time = round(min(hz.times), 3)
                hz.max_time = round(max(hz.times), 3)
                n = len(hz.times)
                mean = sum(hz.times) / n
                hz.rate = 1.0 / mean if mean > 0.0 else 0
                hz.std_dev_time = round(
                    math.sqrt(sum((x - mean) ** 2 for x in hz.times) / n), 4
                )
                print(
                    f"{hz.name:<10} hz:{hz.hz} rate:{hz.rate:.2f} min:{hz.min_time:.3f} max:{hz.max_time:.3f} accum_dev:{hz.accum_dev_time:.3f} std dev:{hz.std_dev_time:.3f}"
                )
                self.csv_record(hz)
            hz.new_start_time = cur_time
            hz.hz = 0
            hz.min_time = 0.0
            hz.max_time = 0.0
            hz.accum_dev_time = 0.0
            hz.times = []
        else:
            hz.hz += 1
            hz.accum_dev_time += abs(hz.fixed_hz_diff_time * hz.hz - diff_time)
            diff_time = cur_time - hz.last_time
            hz.times += [diff_time]

        hz.last_time = cur_time

    def callback_hz_string(self, msg):
        hz = self.hz_string
        self.cal_hz(hz)

    def callback_hz_float(self, msg):
        hz = self.hz_float
        self.cal_hz(hz)

    def callback_hz_float_array(self, msg):
        hz = self.hz_float_array
        self.cal_hz(hz)

    def callback_hz_int(self, msg):
        hz = self.hz_int
        self.cal_hz(hz)

    def callback_hz_int_array(self, msg):
        hz = self.hz_int_array
        self.cal_hz(hz)


def main(args=None):
    rclpy.init(args=args)
    app = topicSubscribeHz(sub_topic_list, Config)
    rclpy.spin(app)
    app.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
